from enum import Enum
from typing import List, Dict, Tuple
import tools
from edge import Edge
import numpy as np
from data import WatchEvent, UserState

class User:
    def __init__(
            self, uid: int, 
            pref: np.ndarray, 
            pick_reco: int,
            videos: Dict[str, Dict],
            user_record: Dict[int, Tuple[int, UserState]],
            upd: float,
            random_select: bool = True) -> None:
        self.uid = uid
        self.pref = pref
        self.state = UserState.IDLE
        self.videos = videos
        self.edge: Edge = None
        self.user_record = user_record
        self.random_select = random_select
        self.pick_reco = pick_reco
        self.upd = upd
        self._watch_list = None
        self._watched_mask = None
        self._decay = 0.9
        self._vid = None
        self._vdx = 1
        self._pv = None
        self._watch_from = None
        self.trace = [self.state]
    
    def __str__(self) -> str:
        return f'User {self.uid}'

    def __repr__(self) -> str:
        return f'User {self.uid}'

    def get_reco_list(self) -> List[str]:
        return self.edge.get_reco_list(self.uid, self.upd)
    
    def get_srch_list(self) -> List[str]:
        unwatched = np.where([not self.has_watched(i) for i in range(len(self.videos['feature']))])[0]
        random_select = np.random.choice(unwatched, size=min(len(unwatched), 100), replace=False)
        # 按照用户偏好排序
        return sorted(random_select, key=lambda x: tools.cosine_similarity(self.pref, self.videos['feature'][x]), reverse=True)[:10]

    def has_watched(self, vid: int) -> bool:
        return vid in self.user_record

    def set_edge(self, edge) -> None:
        self.edge = edge

    def step(self) -> UserState:
        match self.state:
            case UserState.IDLE:
                pr = self.pick_reco
                self.state = np.random.choice([UserState.RECO, UserState.SRCH], p=[pr, 1-pr], size=1)[0]
            case UserState.RECO:
                self._watch_list = np.array(self.get_reco_list())
                self._watched_mask = np.array([self.has_watched(vid) for vid in self._watch_list])
                self._watch_from = UserState.RECO
                self.state = UserState.SELECT
            case UserState.SRCH:
                self._watch_list = np.array(self.get_srch_list())
                self._watched_mask = np.array([self.has_watched(vid) for vid in self._watch_list])
                self._watch_from = UserState.SRCH
                self.state = UserState.SELECT
            case UserState.SELECT:
                # 以某种方式选择一个视频(用户偏好概率为p), 则用户将以p的概率选择观看该视频
                # 否则继续获取推荐列表
                unwatched = self._watch_list[np.where(self._watched_mask == False)[0]]
                # print(f'User {self.uid} unwatched: {unwatched}')  # debug
                self._vid = None if len(unwatched) == 0 else (np.random.choice(unwatched, size=1)[0] if self.random_select else unwatched[0])
                p = 0 if len(unwatched) == 0 else tools.cosine_similarity(self.pref, self.videos['feature'][self._vid])[0]
                self._pv = p
                self.state = np.random.choice([UserState.WATCH, UserState.IDLE], p=[p, 1-p], size=1)[0]
            case UserState.WATCH:
                self.edge.notify(WatchEvent(self.uid, self._vid, self._vdx))
                self.user_record[self._vid] = (self._vdx, self._watch_from)
                # 若当前用户对视频片段的兴趣为p, 则在下一个时间片段,
                # 其转移到STOP的概率为1-p, 转移到WATCH的概率为p
                idx = np.where(self._watch_list == self._vid)[0][0]
                self._watched_mask[idx] = True
                vlen = self.videos['duration'][self._vid]
                self._vdx += 1
                p = self._pv * self._decay ** self._vdx if self._vdx <= vlen else 0
                self.state = np.random.choice([UserState.WATCH, UserState.PAUSE], p=[p, 1-p], size=1)[0]
            case UserState.PAUSE:
                self._vdx = 1
                self.state = np.random.choice([UserState.PAUSE, UserState.SELECT], p=[0.5, 0.5], size=1)[0]
            case UserState.EXIT:
                self.state = UserState.EXIT
        self.trace.append(self.state)

    def clear(self) -> None:
        self.state = UserState.IDLE
        self._watch_list = None
        self._watched_mask = None
        self._vid = None
        self._vdx = 1
        self._pv = None
        self._watch_from = None
        self.trace = [self.state]

    def satisfaction(self) -> float:
        s, c = 0, 0
        for vid, (vdx, vfm) in self.user_record.items():
            if vfm == UserState.RECO:
                s += tools.cosine_similarity(self.pref, self.videos['feature'][vid])[0]
                c += 1
        return 0 if c == 0 else s / c